#test.py
#!/usr/bin/env python3

""" test neuron network performace
print top1 and top5 err on test dataset
of a model

author baiyu
"""

import argparse

from matplotlib import pyplot as plt
import timm
import torch
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from lora1 import *
from conf import settings
from utils import get_test_dataloader
import re

if __name__ == '__main__':

    parser = argparse.ArgumentParser()
    parser.add_argument('-weights', type=str, required=True, help='the weights file you want to test')
    args = parser.parse_args()
    assert 'lora' in args.weights

    net = timm.create_model(args.weights.split('-')[0], pretrained=False, num_classes=100)
    warp_with_lora(net, 20)
    net.load_state_dict(torch.load(args.weights))
    collapse_lora(net)
    result = re.sub(r'lora', "inter_lora", args.weights)
    torch.save(net.cpu().state_dict(), result)